#!/usr/bin/env python
import roslib; roslib.load_manifest('stack_objects')
import rospy, math, copy
from geometry_msgs.msg import PointStamped, Pose, PoseStamped
from object_manipulation_msgs.msg import Grasp, GraspableObject
from pr2_tasks.tasks import Tasks
from pr2_python import base, gripper
from pr2_python.arm_mover import ArmMover
from pr2_tasks.pickplace import PickPlace
from pr2_tasks.pickplace_definitions import PlaceGoal, PickupGoal
from pr2_tasks.tableware_detection import TablewareDetection
from pr2_python import geometry_tools, transform_listener
from pr2_python.world_interface import WorldInterface
from pr2_python.planning_scene_interface import get_planning_scene_interface
from object_manipulation_msgs.srv import FindClusterBoundingBox
from Room3 import *

FIND_CLUSTER_BOUNDING_BOX_SERVICE = '/find_cluster_bounding_box'

class ObjectUnstacker:

   def detect_objects(self):
      det = TablewareDetection()
      head_point = PointStamped()
      head_point.header.frame_id = '/torso_lift_link'
      head_point.point.x = 0.6
      res = det.detect_objects(point_head_at=head_point)
      rospy.loginfo('Detected objects with the following labels and poses:')
      return res

   def separate_graspable_objects(self, res):    
      global wi  
      for pg in res.pickup_goals:
         print pg.label
         if pg.label is 'graspable':
            y_max = pg.target.cluster.points[0].y
            y_min = pg.target.cluster.points[0].y
            for point in pg.target.cluster.points:
               if point.y > y_max:
                  y_max = point.y
               if point.y < y_min:
                  y_min = point.y
            print y_max - y_min
            if (y_max - y_min < 0.14):
               pg.label = 'cups'
               print "Found stack of cups"
            elif (y_max - y_min < 0.2):
               pg.label = 'bowls'
               print "Found stack of bowls"
            else:
               pg.label = 'plates'
               print "Found stack of plates"
      for pg_cups in res.pickup_goals:
         if pg_cups.label is 'cups':
            bounding_box_cups = self.getClusterBoundingBox(pg_cups)
            for pg in res.pickup_goals:
               if pg.collision_object_name is not pg_cups.collision_object_name:
                  bounding_box_object = self.getClusterBoundingBox(pg)
                  if geometry_tools.near(bounding_box_cups.pose.pose, bounding_box_object.pose.pose, dist_eps = max([bounding_box_cups.box_dims.x, bounding_box_cups.box_dims.y, bounding_box_cups.box_dims.z]), angle_eps=3.14159):
                     wi.remove_collision_object(pg.collision_object_name)
                     cups_collision_object_name = pg_cups.collision_object_name
                     pg_cups.target.cluster.points.append(pg.target.cluster.points)
                     pg_cups.target.cluster.channels.append(pg.target.cluster.channels)
                     co = wi.collision_object(pg_cups.collision_object_name)
                     wi.remove_collision_object(pg_cups.collision_object_name)
                     new_bb = self.getClusterBoundingBox(pg)
                     co.poses = [new_bb.pose.pose]
                     wi.add_object(co)
                     rospy.loginfo("We successfully merged multiple clusters")                     
                     return self.getClusterBoundingBox(pg)
            return bounding_box_cups          
 
   def pickup_cups(self, pg, bres):
      tasks = Tasks()
      pickplace = PickPlace()
      alt_arm_name = ''
      if bres.pose.pose.position.y > 0:
         main_arm_name = 'left_arm'
         alt_arm_name = 'right_arm'            
      else:
         main_arm_name = 'right_arm'
         alt_arm_name = 'left_arm'
   
      grasps = []
      rospy.loginfo("Cups located at pose:")
      rospy.loginfo(bres.pose.pose)
      stack_cup_poses = self.returnStackCupPoses(bres.pose.pose, bres.box_dims)
      for pose in stack_cup_poses:
         grasps.append(self.returnDesiredGraspFromPose(pose, main_arm_name, False))
   
      pg.desired_grasps = grasps
      target_frame_id = pg.target.reference_frame_id
      pg.target = GraspableObject()
      pg.target.reference_frame_id = target_frame_id
      pg.arm_name = main_arm_name
      rospy.loginfo("Picking up with " + pg.arm_name)
      pickplace.pickup(pg)

      left_arm_mover = ArmMover('left_arm')
      right_arm_mover = ArmMover('right_arm')
      cup_pose = self.returnMoveCupPose(main_arm_name)
      reached_goal = False
      if main_arm_name is 'right_arm':
         reached_goal = right_arm_mover.move_to_goal(cup_pose)
         if not reached_goal:
            rospy.loginfo(right_arm_mover.get_exceptions())
      else:
         reached_goal = left_arm_mover.move_to_goal(cup_pose)  
         if not reached_goal:
            rospy.loginfo(left_arm_mover.get_exceptions())
      self.reset_planning_scene()
      b = base.Base()
      b.move_to(-0.76, 0.32, -3.086)
      return pg    

   def unstack_cups(self, pg, bres): 
      pickplace = PickPlace() 
      global wi
      main_arm_name = pg.arm_name
      if pg.arm_name is 'right_arm':
         alt_arm_name = 'left_arm'
      else:
         alt_arm_name = 'right_arm'

      left_arm_mover = ArmMover('left_arm')
      right_arm_mover = ArmMover('right_arm')
      cup_pose = self.returnCupPose(main_arm_name)
      reached_goal = False
      if main_arm_name is 'right_arm':
         reached_goal = right_arm_mover.move_to_goal(cup_pose)
         if not reached_goal:
            rospy.loginfo(right_arm_mover.get_exceptions())
      else:
         reached_goal = left_arm_mover.move_to_goal(cup_pose)  
         if not reached_goal:
            rospy.loginfo(left_arm_mover.get_exceptions())
      
      stack_cups_collision_object = wi.attached_collision_object(pg.collision_object_name)
      grasps = []
      top_cup_poses = self.returnTopObjectPoses(stack_cups_collision_object.object, alt_arm_name)
      for pose in top_cup_poses:
         grasps.append(self.returnDesiredGraspFromPose(pose, alt_arm_name, True))

      go = GraspableObject()
      go.reference_frame_id = pg.target.reference_frame_id
      pg.arm_name = alt_arm_name
      rospy.loginfo("Picking up with " + pg.arm_name)
      pg.desired_grasps = grasps
      pickplace.pickup(pg)
      pg.arm_name = main_arm_name
      return pg
            
   def placeCups(self, pg, bres):
      main_arm_name = pg.arm_name
      if pg.arm_name is 'right_arm':
         alt_arm_name = 'left_arm'
      else:
         alt_arm_name = 'right_arm'
      pickplace = PickPlace()
      b = base.Base()
      b.move_manipulable_pose(-1.15, 0.25, 1, group='torso', try_hard=True)
      top_cup_place = copy.deepcopy(bres.pose)
      top_cup_place.pose.position.y += 0.2
      top_cup_place.pose.position.z += bres.box_dims[2]/2 + 0.18
      top_cup_place.pose.orientation = geometry_tools.euler_to_quaternion(0, 3.14159/2, 0)

      b.move_to(-0.76, 0.32, -3.086)
      place_pose = transform_listener.transform_pose_stamped('/map', bres.pose, use_most_recent=True)
      #b.move_manipulable_pose(place_pose.pose.position.x, place_pose.pose.position.y, place_pose.pose.position.z)
      #place_pose.pose.position.x += 0.1

      stack_place_goal = PlaceGoal(main_arm_name, [place_pose],
                       collision_support_surface_name = pg.collision_support_surface_name,
                       collision_object_name = pg.collision_object_name)
      
      try:
         pickplace.place(stack_place_goal)
      except:
         rospy.loginfo("Failed to place the stack of cups")

      self.reset_planning_scene()
      self.move_arms_back()
      place_pose = transform_listener.transform_pose_stamped('/map', top_cup_place, use_most_recent=True)
      b.move_manipulable_pose(place_pose.pose.position.x, place_pose.pose.position.y, place_pose.pose.position.z)

      top_cup_goal = PlaceGoal(alt_arm_name, [place_pose],
                       collision_support_surface_name = pg.collision_support_surface_name,
                       collision_object_name = 'top_cup')
      try:
         pickplace.place(top_cup_goal)
      except:
         rospy.loginfo("Failed to place the stack of cups")

   def returnDesiredGraspFromPose(self, pose, arm_name, top_cup):
      desired_grasp = Grasp()
      desired_grasp.grasp_pose = pose
      if arm_name is 'right_arm':
         desired_grasp.pre_grasp_posture.name.append("r_gripper_l_finger_joint")
         desired_grasp.pre_grasp_posture.name.append("r_gripper_r_finger_joint")
         desired_grasp.pre_grasp_posture.name.append("r_gripper_l_finger_tip_joint")
         desired_grasp.pre_grasp_posture.name.append("r_gripper_r_finger_tip_joint")
         desired_grasp.grasp_posture.name.append("r_gripper_l_finger_joint")
         desired_grasp.grasp_posture.name.append("r_gripper_r_finger_joint")
         desired_grasp.grasp_posture.name.append("r_gripper_l_finger_tip_joint")
         desired_grasp.grasp_posture.name.append("r_gripper_r_finger_tip_joint")
      if arm_name is 'left_arm':
         desired_grasp.pre_grasp_posture.name.append("l_gripper_l_finger_joint")
         desired_grasp.pre_grasp_posture.name.append("l_gripper_r_finger_joint")
         desired_grasp.pre_grasp_posture.name.append("l_gripper_l_finger_tip_joint")
         desired_grasp.pre_grasp_posture.name.append("l_gripper_r_finger_tip_joint")
         desired_grasp.grasp_posture.name.append("l_gripper_l_finger_joint")
         desired_grasp.grasp_posture.name.append("l_gripper_r_finger_joint")
         desired_grasp.grasp_posture.name.append("l_gripper_l_finger_tip_joint")
         desired_grasp.grasp_posture.name.append("l_gripper_r_finger_tip_joint")
      if top_cup:
         desired_grasp.pre_grasp_posture.position = [0.3, 0.3, 0.3, 0.3]
      else:
         desired_grasp.pre_grasp_posture.position = [0.5, 0.5, 0.5, 0.5]
      desired_grasp.pre_grasp_posture.effort = [100.0, 100.0, 100.0, 100.0]
      desired_grasp.grasp_posture.position = [0.0, 0.0, 0.0, 0.0]
      desired_grasp.grasp_posture.effort = [100.0, 100.0, 100.0, 100.0]
      desired_grasp.success_probability = 1
      desired_grasp.min_approach_distance = 0.05
      desired_grasp.desired_approach_distance = 0.5
      return desired_grasp

   def returnTopObjectPoses(self, collision_object, arm_name):
      poses = []
      for shape in collision_object.shapes:
         rospy.loginfo("Shape type: " + str(shape.type))
         if shape.type is 1:
            pose = transform_listener.transform_pose('/base_link', collision_object.header.frame_id, collision_object.poses[0], use_most_recent=True)
            pose.position.z += shape.dimensions[2]/2 + 0.16
            for i in range(8):
               angle = i*2*3.14159/8
               rospy.loginfo("Shape.dimensions " + str(shape.dimensions[1]))
               pose.position.x += math.sin(angle) * shape.dimensions[1]/2
               pose.position.y += math.cos(angle) * shape.dimensions[1]/2
               pose.orientation = geometry_tools.euler_to_quaternion(0, 3.14159/2, -angle)
               rospy.loginfo("Top cup pose: ")
               rospy.loginfo(pose)
            poses.append(pose)
      return poses

   def returnStackCupPoses(self, bb_pose, bb_dims):
      poses = []
      for i in range(24):
         layer = math.floor(i/8)
         angle = i*2*3.14159/8
         pose = copy.deepcopy(bb_pose)
         pose.position.x -= math.cos(angle) * math.cos(3.14159/6) * 0.165
         pose.position.y -= math.sin(angle) * math.cos(3.14159/6) * 0.165+0.02	
         pose.position.z += math.sin(3.14159/6)*0.18 - (bb_dims.z/2.5 - layer * (bb_dims.z/6))
         pose.orientation = geometry_tools.euler_to_quaternion(0, 3.14159/6, angle)
         poses.append(pose)
      return poses

   def returnMoveCupPose(self, arm_name):
      cup_pose = PoseStamped()
      cup_pose.pose.position.x = 0.5
      cup_pose.pose.position.z -= 0
      cup_pose.header.frame_id = '/torso_lift_link'
      if arm_name is 'right_arm':
         cup_pose.pose.orientation = geometry_tools.euler_to_quaternion(0, 3.14159/6, 3.14159/2)
      if arm_name is 'left_arm':
         cup_pose.pose.orientation = geometry_tools.euler_to_quaternion(0, 3.14159/6, -3.14159/2)
      return cup_pose

   def returnCupPose(self, arm_name):
      cup_pose = PoseStamped()
      cup_pose.pose.position.x = 0.5
      cup_pose.pose.position.z -= 0.15
      cup_pose.header.frame_id = '/torso_lift_link'
      if arm_name is 'right_arm':
         cup_pose.pose.orientation = geometry_tools.euler_to_quaternion(0, 3.14159/6, 3.14159/2)
      if arm_name is 'left_arm':
         cup_pose.pose.orientation = geometry_tools.euler_to_quaternion(0, 3.14159/6, -3.14159/2)
      return cup_pose

   def move_arms_back(self):
      tasks = Tasks()
      tasks.move_arms_to_side()

   def reset_everything(self):
      global wi
      wi = WorldInterface()
      wi.reset()
      psi = get_planning_scene_interface()
      psi.reset()
      rospy.sleep(2.0)

   def reset_planning_scene(self):
      psi = get_planning_scene_interface()
      psi.reset()
      rospy.sleep(2.0)


   def getClusterBoundingBox(self, pg):
      bbox_srv = rospy.ServiceProxy(FIND_CLUSTER_BOUNDING_BOX_SERVICE, FindClusterBoundingBox)
      rospy.loginfo('Waiting for bounding box service to fill pickup goal.')
      bbox_srv.wait_for_service()
      rospy.loginfo('Found service')
      return bbox_srv(cluster=pg.target.cluster)
      
   def unstack_objects(self):
      rospy.loginfo("Moving arms to starting position")
      self.move_arms_back()
      rospy.loginfo("Resetting collision models")
      self.reset_everything()
      rospy.loginfo("Moving to table to pick up cups")
      room3 = Room3()
      room3.goto_dirty_table()
      rospy.loginfo("Detecting tableware objects")
      res = self.detect_objects()
      rospy.loginfo("Determining type of graspable objects")
      bres = self.separate_graspable_objects(res)
      for pg in res.pickup_goals:
         if pg.label is 'cups':
            pg = self.pickup_cups(pg, bres)
            pg = self.unstack_cups(pg, bres)
            room3.goto_table()
            #self.move_arms_back()
            self.placeCups(pg, bres)
            return

